//
// Created by vivi on 07/02/2018.
//

#pragma once

#include <algorithm>
#include "core/precompiled/precompiled.h"
#include <Eigen/Eigen>
#include <istream>

namespace HairEngine {
	class StrandBasedHairGeometry {
	Public:
		typedef StrandBasedHairGeometry *Ptr;

		struct StrandParticles {
			const Eigen::Vector3f *particles;
			size_t size;
		};

	Public:
		/*
		 * We can create a StrandBasedHairGeometry from a .hair file
		 */
		explicit StrandBasedHairGeometry(const std::string & filePath, const Eigen::Affine3f & affine = Eigen::Affine3f::Identity());
		explicit StrandBasedHairGeometry(std::istream & is, const Eigen::Affine3f & affine = Eigen::Affine3f::Identity());
		~StrandBasedHairGeometry();

		StrandBasedHairGeometry(const StrandBasedHairGeometry & other); //Copy constructor
		StrandBasedHairGeometry(StrandBasedHairGeometry && other); //Move constructor

		/*
		 * This constructor can be used to initialized a StrandBasedHairGeometry from code.
		 * a ParticleIterator instance p must have proper type (*p) could be assigned to a Eigen::Vector3f,
		 * and a StrandCountIterator instance s must have proper type (*s) could be assigned a size_t.
		 */
		template <typename ParticleIterator, typename StrandCountIterator>
		StrandBasedHairGeometry(const ParticleIterator & particleBegin,
		                        const ParticleIterator & particleEnd,
		                        const StrandCountIterator & strandCountBegin,
		                        const StrandCountIterator & strandCountEnd) {
			std::vector<Eigen::Vector3f> particlesVector;
			std::vector<size_t> endIndexVector;

			for (auto pIt = particleBegin; pIt != particleEnd; ++pIt)
				particlesVector.push_back(*pIt);

			for (auto sIt = strandCountBegin; sIt != strandCountEnd; ++sIt) {
				size_t lastStrandCount = (!endIndexVector.empty() ? endIndexVector.back() : 0);
				endIndexVector.push_back(lastStrandCount + (*sIt));
			}

			//Move the vector to the pointer
			particles = new Eigen::Vector3f[particlesVector.size()];
			endIndex = new size_t[endIndexVector.size()];
			nparticle = particlesVector.size();
			nstrand = endIndexVector.size();

			std::copy(particlesVector.begin(), particlesVector.end(), particles);
			std::copy(endIndexVector.begin(), endIndexVector.end(), endIndex);
		};

		size_t nstrand; //Number of strands in the hair geometry
		size_t nparticle; //Number of the particle in the hair geometry

		inline const Eigen::Vector3f *particleArray() const {
			return particleArrayOnStrand(0);
		}

		inline const Eigen::Vector3f *particleArrayOnStrand(size_t i) const {
			Assert(i >= 0 && i < nstrand);
			return particles + (i != 0 ? endIndex[i - 1] : 0);
		}

		inline size_t particleSizeOnStrand(size_t i) const {
			Assert(i >= 0 && i < nstrand);
			return endIndex[i] - (i != 0 ? endIndex[i - 1] : 0);
		}

		inline const StrandParticles strandStructure(size_t i) const {
			return { particleArrayOnStrand(i), particleSizeOnStrand(i) };
		}

		void writeToFile(const std::string & filePath) const; //Write all the geometry to the filePath

		/*
		 * Sometimes we might want to change all particle positions.
		 * A typical case is a HairSimulator might complete simulation after a timestep,
		 * and tries to write the simulation result into a file. First, it might use this
		 * method to assign particles' positions. And then use "writeToFile" to write the hair
		 * geometry to the specific file
		 */
		template <typename Vector3fArrayIterator>
		void assignParticlePositions(const Vector3fArrayIterator & begin, const Vector3fArrayIterator & end) {
			const Eigen::Vector3f *lastCopyParticlePosition = std::copy(begin, end, particles);
			Assert(lastCopyParticlePosition - particles == nparticle);
		}

		/*
		 * Resample the hair geometry and create a new one, the sampleStep parameters
		 * indicates that how the step should take for sampling. For example, if the sampleStep is 10, we will only
		 * get 1/10 strands for the new hair geometry.
		 */
		StrandBasedHairGeometry resample(size_t sampleStep) {
			auto ret = StrandBasedHairGeometry();

			ret.nstrand = std::max((size_t)1, nstrand / sampleStep);
			ret.nparticle = endIndex[ret.nstrand - 1];

			ret.particles = new Eigen::Vector3f[ret.nparticle];
			ret.endIndex = new size_t[ret.nstrand];

			std::copy(particles, particles + ret.nparticle, ret.particles);
			std::copy(endIndex, endIndex + ret.nstrand, ret.endIndex);

			return ret;
		}

		template <typename StrandIndexIterator>
		StrandBasedHairGeometry resample(const StrandIndexIterator & begin, const StrandIndexIterator & end) {
			StrandBasedHairGeometry ret;

			ret.nstrand = 0;
			ret.nparticle = 0;
			for (auto it = begin; it != end; ++it) {
				++ret.nstrand;
				ret.nparticle += particleSizeOnStrand(*it);
			}

			ret.particles = new Eigen::Vector3f[ret.nparticle];
			ret.endIndex = new size_t[ret.nstrand];

			auto particlePtr = ret.particles;
			auto endIndexPtr = ret.endIndex;
			for (auto it = begin; it != end; ++it) {

				const size_t i = *it;

				auto currentSize = particleSizeOnStrand(i);
				auto currentParticlePtr = particleArrayOnStrand(i);

				std::copy(currentParticlePtr, currentParticlePtr + currentSize, particlePtr);
				*endIndexPtr = (endIndexPtr != ret.endIndex ? *(endIndexPtr - 1) : 0) + currentSize;

				particlePtr += currentSize;
				++endIndexPtr;
			}

			return ret;
		}

	Protected:

		void loadFromStream(std::istream & is, const Eigen::Affine3f & affine);

		StrandBasedHairGeometry() = default; //A default constructor used inside

		Eigen::Vector3f *particles = nullptr; //The particles storage arrays, the size if equal to the particle size

		/*
		 * The end index (not including) for a strand, suppose we want the particles in strand 1, we could get it from
		 * particles[endIndex[0]] to particles[endIndex[1] - 1].
		 * The size of endIndex is equal to the number of strand
		 */
		size_t *endIndex = nullptr;
	};
}
